#include "StdAfx.h"

#include <iostream>
#include <fstream>
#include <Windows.h>
#include <psapi.h>

#define PSAPI_VERSION 1
using namespace std;

#ifndef MAKEULONGLONG
#define MAKEULONGLONG(ldw, hdw) ((ULONGLONG(hdw) << 32) | ((ldw) & 0xFFFFFFFF))
#endif
  
#ifndef MAXULONGLONG
#define MAXULONGLONG ((ULONGLONG)~((ULONGLONG)0))
#endif

HANDLE CProcess::hProcessSnap = NULL;
PROCESSENTRY32 CProcess::pe32 = {0};
DWORD CProcess::pidFind = 0;
TCHAR CProcess::szProcessFind[MAX_PATH] = L"";

//-----------------------------------------------------------------------------

CProcess::CProcess(PROCESSENTRY32 pe)
{	
	wcscpy_s(this->szProcessName, 256, pe.szExeFile);	
	dwProcessId = pe.th32ProcessID;
	hProcess = GetHandle(dwProcessId);
	
	EnableDebugPrivilege(hProcess);

	if(hProcess != NULL) {		
	
		MODULEENTRY32 mainMod = GetModuleInfo(szProcessName);
		dwBaseAddress = (DWORD)(ULONG_PTR)mainMod.modBaseAddr;				
		wcscpy_s(szExePath, MAX_PATH, mainMod.szExePath);		
		if(dwBaseAddress == 0) {
			HANDLE hThread = OpenThread(THREAD_QUERY_INFORMATION | THREAD_GET_CONTEXT, TRUE, GetMainThread().th32ThreadID);
			
			CONTEXT context;
			context.ContextFlags = CONTEXT_FULL;
			GetThreadContext(hThread, &context);			
			BYTE buffer[10];
			Read(context.Ebx + 8, buffer, sizeof(DWORD));
			dwBaseAddress = *(DWORD*)buffer;			
		}

		// Find PEOffset
		BYTE buffer[4];
		Read(0x3C, buffer, 4);
		DWORD dwPEOffset = *(DWORD*)buffer;

		// Read Data Base Address
		Read(dwPEOffset + 0x30, buffer, 4);
		dwDataAddress = dwBaseAddress +  *(DWORD*)buffer;

		// Data Code Address
		dwCodeAddress = dwBaseAddress + 0x1000;
	}
}

CProcess::~CProcess(void)
{
}

void CProcess::EnableDebugPrivilege(HANDLE hProcess)
{
    HANDLE hToken;
    TOKEN_PRIVILEGES tokenPriv;
    LUID luidDebug;
    if(OpenProcessToken(hProcess, TOKEN_ADJUST_PRIVILEGES, &hToken) != FALSE) {
        if(LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &luidDebug) != FALSE)
        {
            tokenPriv.PrivilegeCount           = 1;
            tokenPriv.Privileges[0].Luid       = luidDebug;
            tokenPriv.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
            AdjustTokenPrivileges(hToken, FALSE, &tokenPriv, sizeof(tokenPriv), NULL, NULL);
        }
    }
}

CProcess* CProcess::FindProcess(DWORD pid) {
	return FindProcess(pid, L"");
}

CProcess* CProcess::FindProcess(TCHAR* szProcessName) {
	return FindProcess(0, szProcessName);
}

CProcess* CProcess::FindProcess(DWORD pid, TCHAR* szProcessName) {
	CloseHandle(hProcessSnap);
	hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);	
	if(hProcessSnap == INVALID_HANDLE_VALUE) {
		return NULL;
	}

	pe32.dwSize = sizeof(PROCESSENTRY32);

	pidFind = pid;
	wcscpy_s(szProcessFind, MAX_PATH, szProcessName);

	CProcess* process;
	if(Process32First(hProcessSnap, &pe32)) {		
		if((pidFind > 0 && pe32.th32ProcessID == pidFind) || (pidFind == 0 && !_wcsicmp(szProcessFind, pe32.szExeFile))) {				
			process = new CProcess(pe32);
		} else {
			process = FindProcessNext();
		}
	}
	
	return process;
}

CProcess* CProcess::FindProcessNext() {
	while (Process32Next(hProcessSnap, &pe32)) {	
		if((pidFind > 0 && pe32.th32ProcessID == pidFind) || (pidFind == 0 && !_wcsicmp(szProcessFind, pe32.szExeFile))) {				
			CProcess* process = new CProcess(pe32);
			return process;
		}
	} 

	return NULL;
}

HANDLE CProcess::GetHandle( const DWORD dwProcessId_i )
{
   // Get process handle  
	HANDLE ahmProcess = OpenProcess(PROCESS_TERMINATE|PROCESS_CREATE_THREAD|PROCESS_SET_SESSIONID|PROCESS_VM_OPERATION|PROCESS_VM_READ|PROCESS_VM_WRITE|PROCESS_DUP_HANDLE|PROCESS_CREATE_PROCESS|PROCESS_SET_QUOTA|PROCESS_SET_INFORMATION|PROCESS_QUERY_INFORMATION,  TRUE, dwProcessId_i);   
	if( !ahmProcess )
	{
		
		return NULL;
	}

   HANDLE ahmDuplicateProcHandle;   
   if(!DuplicateHandle( GetCurrentProcess(), ahmProcess, GetCurrentProcess(), &ahmDuplicateProcHandle, PROCESS_TERMINATE|PROCESS_CREATE_THREAD|PROCESS_SET_SESSIONID|PROCESS_VM_OPERATION|PROCESS_VM_READ|PROCESS_VM_WRITE|PROCESS_DUP_HANDLE|PROCESS_CREATE_PROCESS|PROCESS_SET_QUOTA|PROCESS_SET_INFORMATION|PROCESS_QUERY_INFORMATION, FALSE, 0 )) {
	   //showLastError();
	   //mi(20);
	   return NULL;
   }   
   
   return ahmDuplicateProcHandle;
}

BOOL CProcess::Write(LPCWSTR szFile) {	
	std::ifstream fi(szFile, std::ifstream::in);
	DWORD_PTR dwOffset;
	char szData[256];
	BYTE buffer[128];
	
	while(fi.good()) {
		fi >> std::hex >> dwOffset;
		fi >> szData;
		if(strlen(szData) > 0) {
			int iLen = u::StringToByteArray(szData, buffer);
			if(!WriteProcessMemory(hProcess, (LPVOID)dwOffset, buffer, iLen, NULL)) {
				return false;
			}
		}
	}
	
	fi.close();

	return true;
}


BOOL CProcess::Verify(LPCWSTR szFile) {	
	std::ifstream fi(szFile, std::ifstream::in);
	if(!fi.is_open()) {		
		return false;
	}
	
	DWORD_PTR dwOffset;
	char szData[256];
	BYTE bfFile[128];
	BYTE bfExe[128];
	
	while(fi.good()) {
		fi >> std::hex >> dwOffset;		
		fi >> szData;
		if(strlen(szData) > 0) {
			int iLen = u::StringToByteArray(szData, bfFile);
			if(!ReadProcessMemory(hProcess, (LPVOID)dwOffset, bfExe, iLen, NULL)) {
				fi.close();
				return false;
			}
			if(memcmp(bfFile, bfExe, iLen)) {
				fi.close();
				return false;
			}
		}
	}
	
	fi.close();
	return true;
}

BOOL CProcess::Write(DWORD_PTR dwOffset, BYTE* buffer, int nByte) {		
	return WriteProcessMemory(hProcess, (LPVOID)dwOffset, buffer, nByte, NULL);	
}

BOOL CProcess::Read(DWORD_PTR dwOffset, BYTE* buffer, int nByte) {	
	return ReadProcessMemory(hProcess, (LPVOID)dwOffset, buffer, nByte, NULL);	
}


MODULEENTRY32 CProcess::GetModuleInfo(TCHAR* szModuleName) {
	 
	MODULEENTRY32 lpModuleEntry = {0};
	memset(&lpModuleEntry, 0, sizeof(lpModuleEntry));

	HMODULE hMods[1024] = {0};
	DWORD cbNeeded;

	TCHAR szModPatern[MAX_PATH];
	wcscpy_s(szModPatern, MAX_PATH, szModuleName);	
	
	if( EnumProcessModules(hProcess, hMods, sizeof(hMods), &cbNeeded))
    {
		MODULEINFO modInfo;
		TCHAR szModName[MAX_PATH];
		
        for (int i = 0; i < (int)(cbNeeded / sizeof(HMODULE)); i++ )
        {
			GetModuleBaseName( hProcess, hMods[i], szModName, sizeof(szModName) / sizeof(TCHAR));						

			if (!_wcsicmp(szModName, szModPatern ))
            {				
                GetModuleInformation(hProcess, hMods[i], &modInfo, sizeof( MODULEINFO ));				
				lpModuleEntry.hModule = hMods[i];
				lpModuleEntry.modBaseAddr = (BYTE*)modInfo.lpBaseOfDll;
				GetModuleFileNameEx(hProcess, lpModuleEntry.hModule, lpModuleEntry.szExePath, MAX_PATH);				
				// DANGEROUS....I HAVENT FILLED ALL OF MODULE ENTRY MEMBERS
				return lpModuleEntry;
            }
        }
    }

	return lpModuleEntry;
}

BOOL CProcess::Terminate(int nExitCode) {
	return TerminateProcess(hProcess, nExitCode);
}


BOOL CProcess::Find(DWORD* dwOffset, BYTE* bPattern, int nBytes) {	
	if(*dwOffset < dwBaseAddress) {
		*dwOffset = dwBaseAddress;
	}

	return Find(dwOffset, dwDataAddress, bPattern, nBytes);
}

BOOL CProcess::Find(DWORD* dwOffset, DWORD dwTo, BYTE* bPattern, int nBytes) {	
	const int nBufferSize = 2048;
	BYTE bSrc[nBufferSize];

	int i;
	for(; *dwOffset < dwTo;) {
		Read(*dwOffset, bSrc, nBufferSize);		
		for(i = 0; i < nBufferSize - nBytes; i++) {
			if(!u::ByteArrayCompare(bPattern, bSrc + i, nBytes)) {				
				*dwOffset += i;
				return TRUE;
			}	
		}
		*dwOffset += i;
	}
	
	return FALSE;
}

THREADENTRY32 CProcess::GetMainThread() 
{
	HANDLE h = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, dwProcessId);
	THREADENTRY32 te = {0};
	THREADENTRY32 teMain = {0};
	ULONGLONG ullMinCreateTime = MAXULONGLONG;

	if (h != INVALID_HANDLE_VALUE) {
		te.dwSize = sizeof(te);
		if (Thread32First(h, &te)) {
			do {
				if (te.th32OwnerProcessID == dwProcessId) {
					CloseHandle(h);
					return te;
				}
				te.dwSize = sizeof(te);
			} while (Thread32Next(h, &te));
		}
		CloseHandle(h);
	}
	return te;
}


BOOL CProcess::Verify(DWORD dwOffset, BYTE* bTest, int nByte) {
	BYTE bRead[256];
	Read(dwOffset, bRead, nByte);
	for(int i = 0; i < nByte; i++) {
		if(bRead[i] != bTest[i]) {
			return false;
		}
	}

	return true;
}